import matplotlib.pyplot as plt
import numpy as np
import copy

path = "./resources/data.txt"
density = []
sweetness = []
data = []


def loadData(path):
    with open(path, 'r') as fp:
        for line in fp.readlines():
            tempList = line.strip().split(',')
            data.append([float(tempList[1]), float(tempList[2])])
    fp.close()


def show(meansVector, clustering):
    global index
    color = ['red', 'yellow', 'blue']
    size = [25, 100]
    index = -1

    plt.title('k-means')
    plt.xlabel('density')
    plt.ylabel('sweetness')
    for i in range(len(meansVector)):
        plt.scatter(meansVector[i][0], meansVector[i][1], marker='+', c='black', s=100)
    for i in range(len(clustering)):
        for j in range(len(clustering[i])):
            for k in range(len(data)):
                if (data[k] == clustering[i][j]).all():
                    if 8 <= k <= 20:
                        index = 0
                    else:
                        index = 1
                    break
            plt.scatter(clustering[i][j][0], clustering[i][j][1], c=color[i], s=size[index])

    plt.show()


def k_means(n_clusters, max_iter):
    length = len(data)
    # 随机选择K个样本作为初始化向量
    u = []
    for _ in range(n_clusters):
        randomIndex = np.random.randint(length)
        u.append(data[randomIndex])
    # 转换成数组，进行L2范数计算距离
    u_array = np.array(u)
    data_array = np.array(data)
    print(np.shape(data_array[0]))

    currentIter = 0
    c = {}
    while currentIter <= max_iter:
        currentIter += 1
        # 将每一簇划分置为空
        for i in range(n_clusters):
            c[i] = []
        # 计算每个样本与均值向量之间的距离，取距离值最小的下标(簇标记)
        for i in range(len(data_array)):
            minDistIndex = np.linalg.norm(data_array[i]-u_array, axis=1).argmin()
            c[minDistIndex].append(data_array[i])
        # 保留原来的均值向量，计算新的均值向量
        temp = copy.deepcopy(u_array)
        for i in range(n_clusters):
            if len(c[i]):
                u_array[i] = np.mean(c[i], axis=0)

        if (temp == u_array).all():
            return u_array, c
    return u_array, c


if __name__ == '__main__':
    loadData(path)
    meansVector, clustering = k_means(3, 4)
    print(data)
    print(clustering)
    print(meansVector)
    show(meansVector, clustering)
